import torch
from torch.utils.data import Dataset, DataLoader, random_split
from torchvision.datasets import CIFAR100
from torchvision import transforms
import pandas as pd
import numpy as np


# 数据预处理（适配CIFAR100图像数据和Texas100/Purchase100表格数据）
class CustomDataset(Dataset):
    def __init__(self, data, labels, transform=None):
        self.data = data
        self.labels = labels
        self.transform = transform

    def __len__(self):
        return len(self.labels)

    def __getitem__(self, idx):
        x = self.data[idx]
        y = self.labels[idx]
        if self.transform:
            x = self.transform(x)
        return x, y


def get_data_loaders(dataset_name, batch_size=512, split_ratio=0.5):
    """
    加载数据集并划分为D_T（教师用）和D_S（学生用）
    dataset_name: 可选"cifar100", "texas100", "purchase100"
    split_ratio: D_T/D_S的划分比例（论文用0.5，同构非重叠）
    """
    # 1. 图像数据预处理（CIFAR100）
    if dataset_name == "cifar100":
        transform = transforms.Compose([
            transforms.Resize((32, 32)),
            transforms.ToTensor(),
            transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
        ])
        # 加载训练集（成员数据）和测试集（非成员数据）
        train_dataset = CIFAR100(root='./datasets/cifar100/', train=True, download=True, transform=transform)
        test_dataset = CIFAR100(root='./datasets/cifar100/', train=False, download=True, transform=transform)

        # 划分D_T（教师训练集）和D_S（学生训练集）：同构非重叠
        train_size = len(train_dataset)
        dt_size = int(train_size * split_ratio)
        ds_size = train_size - dt_size
        dt_dataset, ds_dataset = random_split(train_dataset, [dt_size, ds_size],
                                              generator=torch.Generator().manual_seed(42))

        # 构建DataLoader
        dt_loader = DataLoader(dt_dataset, batch_size=batch_size, shuffle=True, num_workers=4)
        ds_loader = DataLoader(ds_dataset, batch_size=batch_size, shuffle=True, num_workers=4)
        test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False, num_workers=4)
        return dt_loader, ds_loader, test_loader

    # 2. 表格数据预处理（Texas100/Purchase100）
    elif dataset_name in ["texas100", "purchase100"]:
        data_path = f'./datasets/{dataset_name}/'
        train_df = pd.read_csv(f'{data_path}train.csv')
        test_df = pd.read_csv(f'{data_path}test.csv')

        # 提取特征和标签（假设特征列是feature_0~feature_N，标签列是label）
        X_train = train_df.filter(like='feature_').values.astype(np.float32)
        y_train = train_df['label'].values.astype(np.int64)
        X_test = test_df.filter(like='feature_').values.astype(np.float32)
        y_test = test_df['label'].values.astype(np.int64)

        # 转换为Tensor
        X_train = torch.tensor(X_train)
        y_train = torch.tensor(y_train)
        X_test = torch.tensor(X_test)
        y_test = torch.tensor(y_test)

        # 构建Dataset
        train_dataset = CustomDataset(X_train, y_train)
        test_dataset = CustomDataset(X_test, y_test)

        # 划分D_T和D_S
        train_size = len(train_dataset)
        dt_size = int(train_size * split_ratio)
        ds_size = train_size - dt_size
        dt_dataset, ds_dataset = random_split(train_dataset, [dt_size, ds_size],
                                              generator=torch.Generator().manual_seed(42))

        # 构建DataLoader
        dt_loader = DataLoader(dt_dataset, batch_size=batch_size, shuffle=True, num_workers=4)
        ds_loader = DataLoader(ds_dataset, batch_size=batch_size, shuffle=True, num_workers=4)
        test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False, num_workers=4)
        return dt_loader, ds_loader, test_loader

    else:
        raise ValueError("数据集仅支持cifar100、texas100、purchase100")